import argparse
import glob
import re
from itertools import chain
from typing import List

import pandas as pd
from llama_index.core.schema import TextNode
from tqdm import tqdm

from iflytech_assistant.assistant.dataclasses import RagData
from iflytech_assistant.es import index

parser = argparse.ArgumentParser()

parser.add_argument("folder", help="folder path")
parser.add_argument("-i", "--index", help="index name")

args = parser.parse_args()

INDEX_NAME = args.index


def load_mapping():
    target_mapping_df = pd.read_excel(
        ".vscode/生成偏好对应关系表.xlsx", sheet_name="细分场景id"
    )
    tags_mapping_df = pd.read_excel(
        ".vscode/生成偏好对应关系表.xlsx", sheet_name="生成偏好对应关系"
    )
    target_mapping = {}
    for i, row in target_mapping_df.iterrows():
        target_mapping[row["id"]] = row["名称"].replace("发", "").strip()

    tags_mapping = {}

    for i, row in tags_mapping_df.iloc[1:].iterrows():
        for j in range(0, len(row), 2):
            if pd.isna(row.iloc[j]):
                continue
            tags_mapping[row.iloc[j + 1]] = {
                "target": row.keys()[j].replace("发", "").strip(),
                "tags": re.sub(r"[^\u4e00-\u9fa5a-zA-Z]", "", row.iloc[j]),
            }

    return target_mapping, tags_mapping


target_mapping, tags_mapping = load_mapping()


def parse_five_suggestions(content: str) -> List[str]:
    obj = re.match(r"1[.](.*?)2[.](.*?)3[.](.*?)4[.](.*?)5[.](.*?)(A.*|$)", content)
    return [obj.group(i).strip() for i in range(1, 6)]


dfs = []
for file in glob.glob(f"{args.folder}/*.csv"):
    dfs.append(pd.read_csv(file))

for file in glob.glob(f"{args.folder}/*.xlsx"):
    dfs.append(pd.read_excel(file))

# merge all dataframes into one
df = pd.concat(dfs)

# filter out d_submodeid is NaN
df = df.dropna(subset=["d_submodeid"])

# filter out i_clickword is NaN
df = df.dropna(subset=["d_segid"])

result_dict = {
    "userinput": [],
    "d_submodeid": [],
    "用户偏好": [],
    "i_clickword": [],
}
failed = 0
for i, row in tqdm(df.iterrows()):
    userinput = row["userinput"]
    content = row["content"]
    target = row["d_submodeid"]

    if pd.isna(row["createlikecommend"]):
        # tags = [row["d_generateid"]]
        continue
    else:
        tags = row["createlikecommend"].split(",")
    selected = int(row["d_segid"]) - 1

    target_name = target_mapping.get(target, "未知")
    tag_names = [
        tags_mapping.get(tag, {"target": "未知", "tags": "未知"}) for tag in tags
    ]
    try:
        five_suggestions = parse_five_suggestions(content)
    except:
        failed += 1
        continue

    i_clickword = five_suggestions[selected]
    d_submodeid = target_name
    if len(tag_names) > 1:
        try:
            preferenced_style = tag_names[selected]["tags"]
        except IndexError:
            print("index out of range", tag_names, selected, content)
            continue
    else:
        preferenced_style = tag_names[0]["tags"]

    result_dict["userinput"].append(userinput)
    result_dict["d_submodeid"].append(d_submodeid)
    result_dict["用户偏好"].append(preferenced_style)
    result_dict["i_clickword"].append(i_clickword)

print(f"parsed {len(df) - failed} rows")
print(f"failed to parse {failed} rows")
df = pd.DataFrame(result_dict)

df = (
    df.groupby(["d_submodeid", "用户偏好", "userinput"])
    .agg(lambda x: "\n".join(x))
    .reset_index()
)
print(f"after groupby, {len(df)} rows")

nodes: List[TextNode] = []

for i, row in df.iterrows():
    data: RagData = RagData(
        input=row["userinput"],
        target=row["d_submodeid"],
        tag=row["用户偏好"],
        mode="polish",
        examples=row["i_clickword"].split("\n"),
    )
    node: TextNode = data.to_text_node()
    nodes.append(node)
index(nodes, INDEX_NAME)
